
#include "nms_with_mask_custom_tiling.h"
#include "register/op_def_registry.h"
#include "graph/utils/type_utils.h"
#include "tiling/platform/platform_ascendc.h"


namespace optiling {
const uint32_t BLOCK_SIZE = 32;
const uint32_t BUFFER_NUM = 1;
const int32_t DIMS_LIMIT_1 = 2;
const int32_t DIMS_LIMIT_2 = 1;
constexpr int32_t ATTRPOS0 = 0;
constexpr int32_t ATTRPOS1 = 1;
constexpr uint32_t INDEXZERO = 0;
constexpr uint32_t INDEXONE = 1;
constexpr uint32_t CoreNum = 1;
constexpr size_t WS_SYS_SIZE = 0;  // Example workspace size, adjust as needed


static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
    TilingData tiling;
    uint64_t ubSize;
    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ubSize);
    
    auto attrs = context->GetAttrs();
    float iou_threshold = 0.1f;  // 默认值
    float scores_threshold = 0.1f;  // 默认值

    if(attrs) {
        if (attrs->GetFloat(ATTRPOS0)) {
            iou_threshold = *(attrs->GetFloat(ATTRPOS0));
        }
        if (attrs->GetFloat(ATTRPOS1)) {
            scores_threshold = *(attrs->GetFloat(ATTRPOS1));
        }
    }
    // Based on the input length and the number of inputs, the number of bytes of the input data type is obtained
    uint32_t inputNum = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
    if(inputNum / 4 >= 2048) {
        std::cout << "[error]  The max number of box is 2048." << std::endl;
        return ge::GRAPH_FAILED;
    }
    tiling.set_totalLength(inputNum);
    tiling.set_iou_threshold(iou_threshold);
    tiling.set_scores_threshold(scores_threshold);
    
    context->SetBlockDim(CoreNum);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = WS_SYS_SIZE;
    return ge::GRAPH_SUCCESS;
}
}
namespace ge {

} // namespace ge

namespace ops {
class NMSWithMaskCustom : public OpDef {
public:
    explicit NMSWithMaskCustom(const char *name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ ge::DT_FLOAT})
            .Format({ge::FORMAT_ND});
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ ge::DT_FLOAT})
            .Format({ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ ge::DT_UINT8})
            .Format({ge::FORMAT_ND});
        this->Attr("iou_threshold")
            .AttrType(OPTIONAL)
            .Float(0.1f); // 默认值
         this->Attr("scores_threshold")
            .AttrType(OPTIONAL)
            .Float(0.1f); // 默认值

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");
    }
};
OP_ADD(NMSWithMaskCustom);
} // namespace ops
